import copy

def split_conc_cont_columns(df, conc_max_value_num=15, preselected_conc_columns=[], preselected_cont_columns=[]):
    '''
    根据这一列特征的个数判断, 选择出可能是离散x的列。
    must_select里面的列必然被选上，一般是把y放在里面！
    must_drop里面的列必然不被选上
    conc_max_value_num是作为离散变量，最多有那么多个值。多于conc_max_value_num的值被认为是连续变量。

    返回columnsß
    '''
    assert (df.columns.duplicated() == False).all()
    conc_columns = copy.copy(preselected_conc_columns)
    cont_columns = copy.copy(preselected_cont_columns)
    assert len(set(cont_columns) & set(cont_columns)) == 0
    for col in df.columns:
        if (col in conc_columns) or (col in cont_columns):
            continue
        if len(df[col].unique()) <= conc_max_value_num:
            conc_columns.append(col)
        else:
            cont_columns.append(col)
    return conc_columns, cont_columns


def cont_columns(df, must_select=[], must_drop=[], conc_max_value_num=15):
    '''
    根据这一列特征的个数判断, 选择出可能是连续x的列。
    must_select里面的列必然被选上，一般是把y放在里面！
    conc_max_value_num是作为离散变量，最多有那么多个值。多于conc_max_value_num的值被认为是连续变量（和select_conc_columns统一）。

    返回columns
    '''
    assert (df.columns.duplicated() == False).all()
    selected = []
    for col in df.columns:
        if ((len(df[col].unique()) > conc_max_value_num) or (col in must_select)) and (col not in must_drop):
            selected.append(col)
    return selected


def add_prefix_columns(df, prefix, inplace=False):
    assert (df.columns.duplicated() == False).all()
    __replace_dict = {}
    for key in df.columns:
        __replace_dict[key] = prefix + '_' + key
    if inplace:
        df.rename(columns=__replace_dict, inplace=True)
    else:
        return df.rename(columns=__replace_dict, inplace=False)
